# _*_ coding:utf-8 _*_
"""
:File: database.py
:Author: cfp
:Date: 2025/9/10 09:41
:LastEditTime: 2025/9/10 09:41
:LastEditors: cfp
:Description: 
"""
import asyncio
import psycopg2
from psycopg2.extras import RealDictCursor
from contextlib import asynccontextmanager
from application.common.config import Config

config = Config()


class Database:
    def __init__(self):
        self.config = config
        self._pool = None

    def get_connection(self):
        """获取数据库连接"""
        return psycopg2.connect(
            host=self.config.DB_HOST,
            port=self.config.DB_PORT,
            database=self.config.DB_NAME,
            user=self.config.DB_USER,
            password=self.config.DB_PASSWORD,
            cursor_factory=RealDictCursor
        )

    @asynccontextmanager
    async def get_async_connection(self):
        """异步获取数据库连接"""
        conn = None
        try:
            # 在线程池中执行数据库连接
            loop = asyncio.get_event_loop()
            conn = await loop.run_in_executor(None, self.get_connection)
            yield conn
        finally:
            if conn:
                conn.close()

    async def execute_query(self, query, params=None):
        """执行查询"""
        async with self.get_async_connection() as conn:
            cursor = conn.cursor()
            await asyncio.get_event_loop().run_in_executor(
                None, cursor.execute, query, params
            )
            if query.strip().upper().startswith('SELECT'):
                result = await asyncio.get_event_loop().run_in_executor(
                    None, cursor.fetchall
                )
                return result
            else:
                await asyncio.get_event_loop().run_in_executor(
                    None, conn.commit
                )
                return cursor.rowcount

    async def execute_many(self, query, params_list):
        """批量执行"""
        async with self.get_async_connection() as conn:
            cursor = conn.cursor()
            await asyncio.get_event_loop().run_in_executor(
                None, cursor.executemany, query, params_list
            )
            await asyncio.get_event_loop().run_in_executor(
                None, conn.commit
            )
            return cursor.rowcount


# 全局数据库实例
db = Database()